import argparse
from functools import wraps
from neo4j import GraphDatabase
import csv

DEFAULT_HOST = "localhost"
DEFAULT_PORT = "7687"
DEFAULT_USERNAME = "neo4j"
DEFAULT_PASSWORD = "xy001005"

DEFAULT_OUTPUT_PATH = "./"
DEFAULT_OUTPUT_FILENAME = "SortedCriticalNode.csv"


# 用于统计所有关键节点方法加权后的值
final_score_dict = {}
# 按比例缩放每次的score，最大值设置为100
MAX_SCORE=100


def list2dict(func):
    @wraps(func)
    def inner(*args, **kwargs):
        results = func(*args, **kwargs)
        MAX = results[0][2]
        for result in results:
            thisId = result[0]+"_"+result[1]
            final_score_dict[thisId] = final_score_dict.get(thisId, 0) + result[2]*(MAX_SCORE/MAX)
    return inner


@list2dict
def in_degree_analyse(tx):  # 入度数统计
    query = (
        "MATCH ()-[r:depends_on_pkg]->(n:Package) "
        "RETURN n.name, n.version, count(r) as inDegreeCount "
        "ORDER BY inDegreeCount DESC LIMIT 200"
    )
    results = tx.run(query).values('n.name', 'n.version', 'inDegreeCount')
    return results


# 对于没有入度的节点，在关键节点评估上给减10
def zero_in_degree_analyse(tx):  # 没有入度的节点统计
    query = (
        "MATCH (n:Package) "
        "OPTIONAL MATCH ()-[r:depends_on_pkg]->(n:Package) "
        "WITH n, count(r) AS cnt "
        "WHERE cnt = 0 "
        "RETURN n.name, n.version"
    )
    results = tx.run(query).values('n.name', 'n.version')
    for res in results:
        thisId = res[0] + "_" + res[1]
        final_score_dict[thisId] = final_score_dict.get(thisId, 0) - 10


@list2dict
def out_degree_analyse(tx):  # 出度数统计
    query = (
        "MATCH ()<-[r:depends_on_pkg]-(n:Package) "
        "RETURN n.name, n.version, count(r) as outDegreeCount "
        "ORDER BY outDegreeCount DESC"
    )
    results = tx.run(query).values('n.name', 'n.version', 'outDegreeCount')
    return results


@list2dict
def pagerank_analyse(tx):  # 根据PR算法计算
    query = (
        "CALL gds.pageRank.stream('myGraph') "
        "YIELD nodeId, score "
        "RETURN gds.util.asNode(nodeId).name AS name, gds.util.asNode(nodeId).version AS version, score "
        "ORDER BY score DESC"
    )
    results = tx.run(query).values('name', 'version', 'score')
    return results


@list2dict
def closeness_analyse(tx):  # 接近中间性统计
    query = (
        "call gds.alpha.closeness.harmonic.stream('myGraph') "
        "YIELD nodeId, centrality as score "
        "RETURN gds.util.asNode(nodeId).name AS name, gds.util.asNode(nodeId).version AS version, score "
        "ORDER BY score DESC"
    )
    results = tx.run(query).values('name', 'version', 'score')
    return results


@list2dict
def betweenness_analyse(tx):  # 介数中心性统计
    query = (
        "call gds.betweenness.stream('myGraph') "
        "YIELD nodeId, score "
        "RETURN gds.util.asNode(nodeId).name AS name, gds.util.asNode(nodeId).version AS version, score "
        "ORDER BY score DESC"
    )
    results = tx.run(query).values('name', 'version', 'score')
    return results


def build_my_graph(tx, rela_name="depends_on_pkg"):
    query = (
        f"CALL gds.graph.project('myGraph','Package','{rela_name}')"
    )
    tx.run(query)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()  # 创建一个解析对象
    # 向该对象中添加要关注的命令行参数和选项
    parser.add_argument('--host', default=DEFAULT_HOST)
    parser.add_argument('--port', default=DEFAULT_PORT)
    parser.add_argument('--username', default=DEFAULT_USERNAME)
    parser.add_argument('--password', default=DEFAULT_PASSWORD)
    parser.add_argument('--output_path', default=DEFAULT_OUTPUT_PATH, help="存放关键节点排名文件的路径")
    parser.add_argument('--output_file_name', default=DEFAULT_OUTPUT_FILENAME, help="存放关键节点排名文件的文件名")
    parser.add_argument('--analyse_zeroindegree', type=bool, default=False, help="选择计算关键节点时是否无入度节点")
    parser.add_argument('--analyse_indegree', type=bool, default=True, help="选择计算关键节点时是否考虑入度")
    parser.add_argument('--analyse_outdegree', type=bool, default=True, help="选择计算关键节点时是否考虑出度")
    parser.add_argument('--analyse_pagerank', type=bool, default=True, help="选择计算关键节点时是否考虑PageRank算法")
    parser.add_argument('--analyse_closeness', type=bool, default=True, help="选择计算关键节点时是否考虑接近中心性算法")
    parser.add_argument('--analyse_betweenness', type=bool, default=True, help="选择计算关键节点时是否考虑介数中心性算法")
    args = parser.parse_args()  # 进行解析
    # uri(统一资源标识符，表示Web中每一种可用资源位置) = 资源命名机制+存放资源的主机名+端口号
    uri = f'neo4j://{args.host}:{args.port}'
    username = args.username
    password = args.password
    filepath = args.output_path
    filename = args.output_file_name
    # 连接数据库
    driver = GraphDatabase.driver(uri, auth=(username, password),
                                  max_connection_lifetime=3600 * 24 * 30, keep_alive=True)
    with driver.session() as session:
        # 根据包节点和depends关系建立一张投影图以便于查询
        try:
            session.execute_write(build_my_graph)
        except BaseException:  # 已经建立图了
            pass
        if args.analyse_zeroindegree:
            session.execute_read(zero_in_degree_analyse)
        if args.analyse_indegree:
            session.execute_read(in_degree_analyse)
        if args.analyse_outdegree:
            session.execute_read(out_degree_analyse)
        if args.analyse_pagerank:
            session.execute_read(pagerank_analyse)
        if args.analyse_closeness:
            session.execute_read(closeness_analyse)
        if args.analyse_betweenness:
            session.execute_read(betweenness_analyse)
        final_score_list = sorted(final_score_dict.items(), key=lambda x:x[1], reverse=True)
        print("关键软件包排名前50为：")
        for i in range (50):
            print(f'    {final_score_list[i][0]}: {final_score_list[i][1]}')
        with open(filepath + filename, 'w', encoding='utf8', newline='')as fp:
            csv_writer = csv.writer(fp)
            csv_writer.writerow(['Name_Version', 'Score'])
            csv_writer.writerows(final_score_list)
        print(f"所有的关键软件包排名及得分已存储在 {filename}.csv 中")


# See PyCharm help at https://www.jetbrains.com/help/pycharm/
